package com.techzhi.tidb.utils;


import com.techzhi.tidb.model.TableColumnInfo;
import com.techzhi.tidb.pool.TiDBConnectionPoolManager;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import java.sql.*;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * TiDB操作工具类
 * 提供专门的TiDB操作API，封装常用操作
 *
 * @author shouzhi
 * @date 2025-09-01
 */

@Slf4j
public class TiDBOperationUtils {

    @Autowired
    private TiDBConnectionPoolManager connectionPoolManager;

    /**
     * 执行查询
     */
    public List<Map<String, Object>> executeQuery(String sql, Object... params) {
        return executeQuery(sql, 0, params);
    }

    /**
     * 执行查询（带超时）
     */
    public List<Map<String, Object>> executeQuery(String sql, int timeoutSeconds, Object... params) {
        long startTime = System.currentTimeMillis();

        try (Connection conn = connectionPoolManager.getConnection();
             PreparedStatement pstmt = conn.prepareStatement(sql)) {

            if (timeoutSeconds > 0) {
                pstmt.setQueryTimeout(timeoutSeconds);
            }

            setParameters(pstmt, params);

            try (ResultSet rs = pstmt.executeQuery()) {
                List<Map<String, Object>> result = resultSetToList(rs);

                long endTime = System.currentTimeMillis();
                log.debug("TiDB查询执行完成，耗时: {}ms，返回记录数: {}",
                    endTime - startTime, result.size());

                return result;
            }

        } catch (SQLException e) {
            log.error("TiDB查询执行失败，SQL: {}, 参数: {}", sql, params, e);
            throw new RuntimeException("TiDB查询执行失败: " + e.getMessage(), e);
        }
    }

    /**
     * 执行更新（INSERT/UPDATE/DELETE）
     */
    public int executeUpdate(String sql, Object... params) {
        return executeUpdate(sql, 0, params);
    }

    /**
     * 执行更新（带超时）
     */
    public int executeUpdate(String sql, int timeoutSeconds, Object... params) {
        long startTime = System.currentTimeMillis();

        try (Connection conn = connectionPoolManager.getConnection();
             PreparedStatement pstmt = conn.prepareStatement(sql)) {

            if (timeoutSeconds > 0) {
                pstmt.setQueryTimeout(timeoutSeconds);
            }

            setParameters(pstmt, params);
            int affectedRows = pstmt.executeUpdate();

            long endTime = System.currentTimeMillis();
            log.debug("TiDB更新执行完成，耗时: {}ms，影响行数: {}",
                endTime - startTime, affectedRows);

            return affectedRows;

        } catch (SQLException e) {
            log.error("TiDB更新执行失败，SQL: {}, 参数: {}", sql, params, e);
            throw new RuntimeException("TiDB更新执行失败: " + e.getMessage(), e);
        }
    }

    /**
     * 批量执行
     */
    public int[] executeBatch(String sql, List<Object[]> paramsList) {
        return executeBatch(sql, paramsList, 100, 0);
    }

    /**
     * 批量执行（带批量大小和超时控制）
     */
    public int[] executeBatch(String sql, List<Object[]> paramsList, int batchSize, int timeoutSeconds) {
        if (paramsList == null || paramsList.isEmpty()) {
            return new int[0];
        }

        long startTime = System.currentTimeMillis();
        List<Integer> results = new ArrayList<>();

        try (Connection conn = connectionPoolManager.getConnection();
             PreparedStatement pstmt = conn.prepareStatement(sql)) {

            conn.setAutoCommit(false);

            if (timeoutSeconds > 0) {
                pstmt.setQueryTimeout(timeoutSeconds);
            }

            int count = 0;
            for (Object[] params : paramsList) {
                setParameters(pstmt, params);
                pstmt.addBatch();
                count++;

                // 达到批量大小时执行一次
                if (count % batchSize == 0) {
                    int[] batchResults = pstmt.executeBatch();
                    for (int result : batchResults) {
                        results.add(result);
                    }
                    pstmt.clearBatch();
                }
            }

            // 执行剩余的批量操作
            if (count % batchSize != 0) {
                int[] batchResults = pstmt.executeBatch();
                for (int result : batchResults) {
                    results.add(result);
                }
            }

            conn.commit();

            long endTime = System.currentTimeMillis();
            log.debug("TiDB批量执行完成，耗时: {}ms，处理记录数: {}",
                endTime - startTime, paramsList.size());

            return results.stream().mapToInt(Integer::intValue).toArray();

        } catch (SQLException e) {
            log.error("TiDB批量执行失败，SQL: {}, 记录数: {}", sql, paramsList.size(), e);
            throw new RuntimeException("TiDB批量执行失败: " + e.getMessage(), e);
        }
    }

    /**
     * 事务执行
     */
    public <T> T executeInTransaction(TransactionCallback<T> callback) {
        return executeInTransaction(callback, Connection.TRANSACTION_READ_COMMITTED);
    }

    /**
     * 事务执行（指定隔离级别）
     */
    public <T> T executeInTransaction(TransactionCallback<T> callback, int isolationLevel) {
        long startTime = System.currentTimeMillis();

        try (Connection conn = connectionPoolManager.getConnection()) {

            int originalIsolationLevel = conn.getTransactionIsolation();
            boolean originalAutoCommit = conn.getAutoCommit();

            try {
                conn.setTransactionIsolation(isolationLevel);
                conn.setAutoCommit(false);

                T result = callback.doInTransaction(conn);
                conn.commit();

                long endTime = System.currentTimeMillis();
                log.debug("TiDB事务执行完成，耗时: {}ms", endTime - startTime);

                return result;

            } catch (Exception e) {
                try {
                    conn.rollback();
                    log.debug("TiDB事务已回滚");
                } catch (SQLException rollbackEx) {
                    log.error("TiDB事务回滚失败", rollbackEx);
                }
                throw new RuntimeException("TiDB事务执行失败: " + e.getMessage(), e);
            } finally {
                try {
                    conn.setTransactionIsolation(originalIsolationLevel);
                    conn.setAutoCommit(originalAutoCommit);
                } catch (SQLException e) {
                    log.warn("恢复连接原始状态失败", e);
                }
            }

        } catch (SQLException e) {
            log.error("TiDB事务执行失败", e);
            throw new RuntimeException("TiDB事务执行失败: " + e.getMessage(), e);
        }
    }

    /**
     * 获取表结构信息
     */
    public List<TableColumnInfo> getTableStructure(String tableName) {
        String sql = "SELECT " +
                    "COLUMN_NAME, DATA_TYPE, IS_NULLABLE, COLUMN_DEFAULT, COLUMN_COMMENT, " +
                    "CHARACTER_MAXIMUM_LENGTH, NUMERIC_PRECISION, NUMERIC_SCALE, ORDINAL_POSITION, " +
                    "COLUMN_KEY, EXTRA " +
                    "FROM INFORMATION_SCHEMA.COLUMNS " +
                    "WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? " +
                    "ORDER BY ORDINAL_POSITION";

        List<Map<String, Object>> rows = executeQuery(sql,
            connectionPoolManager.getConfigManager().getConfig().getDatabase(), tableName);

        return rows.stream().map(row -> {
            TableColumnInfo info = new TableColumnInfo();
            info.setColumnName((String) row.get("COLUMN_NAME"));
            info.setDataType((String) row.get("DATA_TYPE"));
            info.setNullable("YES".equals(row.get("IS_NULLABLE")));
            info.setDefaultValue((String) row.get("COLUMN_DEFAULT"));
            info.setComment((String) row.get("COLUMN_COMMENT"));

            Object maxLength = row.get("CHARACTER_MAXIMUM_LENGTH");
            if (maxLength != null) {
                info.setColumnSize(((Number) maxLength).longValue());
            }

            Object precision = row.get("NUMERIC_PRECISION");
            if (precision != null) {
                info.setColumnSize(((Number) precision).longValue());
            }

            Object scale = row.get("NUMERIC_SCALE");
            if (scale != null) {
                info.setDecimalDigits(((Number) scale).intValue());
            }

            Object position = row.get("ORDINAL_POSITION");
            if (position != null) {
                info.setOrdinalPosition(((Number) position).intValue());
            }

            info.setPrimaryKey("PRI".equals(row.get("COLUMN_KEY")));
            info.setAutoIncrement("auto_increment".equals(row.get("EXTRA")));

            return info;
        }).collect(Collectors.toList());
    }

    /**
     * 获取所有表名
     */
    public List<String> getAllTableNames() {
        String sql = "SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES " +
                    "WHERE TABLE_SCHEMA = ? AND TABLE_TYPE = 'BASE TABLE' " +
                    "ORDER BY TABLE_NAME";

        List<Map<String, Object>> rows = executeQuery(sql,
            connectionPoolManager.getConfigManager().getConfig().getDatabase());

        return rows.stream()
                  .map(row -> (String) row.get("TABLE_NAME"))
                  .collect(Collectors.toList());
    }

    /**
     * 获取表的行数
     */
    public long getTableRowCount(String tableName) {
        String sql = "SELECT COUNT(*) as row_count FROM " + tableName;
        List<Map<String, Object>> result = executeQuery(sql);
        if (!result.isEmpty()) {
            Object count = result.get(0).get("row_count");
            return ((Number) count).longValue();
        }
        return 0;
    }

    /**
     * 检查表是否存在
     */
    public boolean tableExists(String tableName) {
        String sql = "SELECT COUNT(*) as table_count FROM INFORMATION_SCHEMA.TABLES " +
                    "WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?";

        List<Map<String, Object>> result = executeQuery(sql,
            connectionPoolManager.getConfigManager().getConfig().getDatabase(), tableName);

        if (!result.isEmpty()) {
            Object count = result.get(0).get("table_count");
            return ((Number) count).intValue() > 0;
        }
        return false;
    }

    /**
     * 测试连接
     */
    public boolean testConnection() {
        try {
            List<Map<String, Object>> result = executeQuery("SELECT 1 as test");
            return !result.isEmpty() && "1".equals(String.valueOf(result.get(0).get("test")));
        } catch (Exception e) {
            log.error("TiDB连接测试失败", e);
            return false;
        }
    }

    /**
     * 获取数据库版本信息
     */
    public String getDatabaseVersion() {
        try {
            List<Map<String, Object>> result = executeQuery("SELECT VERSION() as version");
            if (!result.isEmpty()) {
                return (String) result.get(0).get("version");
            }
        } catch (Exception e) {
            log.error("获取TiDB版本信息失败", e);
        }
        return "Unknown";
    }

    // 辅助方法
    private void setParameters(PreparedStatement pstmt, Object... params) throws SQLException {
        if (params != null) {
            for (int i = 0; i < params.length; i++) {
                pstmt.setObject(i + 1, params[i]);
            }
        }
    }

    private List<Map<String, Object>> resultSetToList(ResultSet rs) throws SQLException {
        List<Map<String, Object>> list = new ArrayList<>();
        ResultSetMetaData metaData = rs.getMetaData();
        int columnCount = metaData.getColumnCount();

        while (rs.next()) {
            Map<String, Object> row = new HashMap<>();
            for (int i = 1; i <= columnCount; i++) {
                row.put(metaData.getColumnName(i), rs.getObject(i));
            }
            list.add(row);
        }

        return list;
    }

    /**
     * 事务回调接口
     */
    @FunctionalInterface
    public interface TransactionCallback<T> {
        T doInTransaction(Connection conn) throws Exception;
    }
}
